from datetime import datetime

from langchain.agents import Tool, initialize_agent
from langchain.agents import AgentType
from langchain_ollama import OllamaLLM

# Step 1: 定义一个获取当前时间的函数
def get_current_time(*args, **kwargs):
    """返回当前时间"""
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    message = f"当前时间是：{now}"
    return message

# Step 2: 创建 LangChain 工具
tools = [
    Tool(
        name="GetCurrentTime",
        func=get_current_time,
        description="当你需要知道当前时间时可以使用这个工具",
    )
]

# Step 3: 初始化本地 Ollama 模型
llm = OllamaLLM(model="qwen3:8b")  # 确保你已经在本地拉取了这个模型

# Step 4: 初始化 Agent
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True,
    enable_thinking=False,
    verbose=True,
)


# Step 5: 运行 Agent 来查询时间
prompt = "现在几点了？"
response = llm.invoke(prompt)

print(response)